import numpy as np
import pandas as pd
from sklearn.preprocessing import RobustScaler

class DataPreprocessor:
    def __init__(self, config):
        self.config = config
        self.scaler = RobustScaler()
        
    def handle_missing_values(self, df):
        if self.config.HANDLE_MISSING == 'interpolate':
            # 对每个数值列分别进行插值
            numeric_columns = df.select_dtypes(include=[np.number]).columns
            for col in numeric_columns:
                # 使用更简单的线性插值方法
                df[col] = df[col].interpolate(
                    method='linear',  # 改用线性插值
                    limit_direction='both'
                ).fillna(  # 处理首尾的缺失值
                    method='ffill'
                ).fillna(
                    method='bfill'
                )
                
                # 如果仍有NaN，用列的均值填充
                if df[col].isna().any():
                    df[col] = df[col].fillna(df[col].mean())
                
        elif self.config.HANDLE_MISSING == 'fill':
            # 使用前向填充，然后后向填充
            df = df.fillna(method='ffill').fillna(method='bfill')
            
            # 如果仍有NaN，用列的均值填充
            for col in df.columns:
                if df[col].isna().any():
                    df[col] = df[col].fillna(df[col].mean())
        
        return df
    
    def remove_outliers(self, df):
        for col in self.config.FEATURE_COLUMNS:
            if col not in ['hour', 'day', 'month', 'weekday']:
                z_scores = np.abs((df[col] - df[col].mean()) / df[col].std())
                df.loc[z_scores > self.config.OUTLIER_THRESHOLD, col] = np.nan
        return df
    
    def scale_features(self, df):
        # 分离时间特征
        time_features = df[['hour', 'day', 'month', 'weekday']]
        weather_features = df.drop(['hour', 'day', 'month', 'weekday'], axis=1)
        
        # 只对气象特征进行缩放
        scaled_features = self.scaler.fit_transform(weather_features)
        scaled_df = pd.DataFrame(scaled_features, columns=weather_features.columns)
        
        # 合并回时间特征
        return pd.concat([scaled_df, time_features], axis=1)
    
    def add_time_features(self, df):
        """添加时间特征"""
        # 确保timestamp列存在
        if 'timestamp' not in df.columns:
            raise ValueError("Missing required 'timestamp' column")
        
        # 转换时间戳并添加时间特征
        df['timestamp'] = pd.to_datetime(df['timestamp'])
        df['hour'] = df['timestamp'].dt.hour
        df['day'] = df['timestamp'].dt.day
        df['month'] = df['timestamp'].dt.month
        df['weekday'] = df['timestamp'].dt.weekday
        return df
    
    def process(self, df):
        """修改处理顺序"""
        # 1. 首先添加时间特征
        df = self.add_time_features(df)
        
        # 2. 然后处理异常值
        df = self.remove_outliers(df)
        
        # 3. 处理缺失值
        df = self.handle_missing_values(df)
        
        # 4. 最后进行特征缩放
        df = self.scale_features(df)
        
        return df 